import re
import xml.etree.ElementTree as ET
from loguru import logger

def parse_llm_output(llm_output: str) -> dict:
    """Parses the LLM's XML-like output into a structured dictionary."""
    parsed_data = {
        "think": None,
        "reply": None,
        "clarify": None,
        "tool_calls": [],
        "memory_queries": [],
        "all_plan": None,
        "revise_plan": None
    }

    # Extract <think> tag content first
    think_match = re.search(r"<think>(.*?)</think>", llm_output, re.DOTALL)
    if think_match:
        parsed_data["think"] = think_match.group(1).strip()

    # If a final answer is given, parse it and return immediately
    reply_match = re.search(r"<reply>(.*?)</reply>", llm_output, re.DOTALL)
    if reply_match:
        parsed_data["reply"] = reply_match.group(1).strip()
        return parsed_data

    # If a clarifying question is asked, parse it and return
    clarify_match = re.search(r"<clarify>(.*?)</clarify>", llm_output, re.DOTALL)
    if clarify_match:
        parsed_data["clarify"] = clarify_match.group(1).strip()
        return parsed_data

    # Parse plans
    all_plan_match = re.search(r"<all_plan>(.*?)</all_plan>", llm_output, re.DOTALL)
    if all_plan_match:
        parsed_data["all_plan"] = all_plan_match.group(1).strip()

    revise_plan_match = re.search(r"<revise_plan>(.*?)</revise_plan>", llm_output, re.DOTALL)
    if revise_plan_match:
        parsed_data["revise_plan"] = revise_plan_match.group(1).strip()

    # Parse tool calls and memory queries using XML parser for robustness
    try:
        # Find tool_calls
        tool_calls_match = re.search(r"<tool_calls>(.*?)</tool_calls>", llm_output, re.DOTALL)
        if tool_calls_match:
            # Wrap in a root tag to ensure it's valid XML
            tool_calls_xml = f"<root>{tool_calls_match.group(1)}</root>"
            action_root = ET.fromstring(tool_calls_xml)
            for tool_elem in action_root.findall('tool'):
                tool_call = {
                    "tool_name": tool_elem.findtext("tool_name", "").strip(),
                    "params": {},#ET.tostring(tool_elem.find("params"), encoding='unicode').strip() if tool_elem.find("params") is not None else ""
                    "expect": tool_elem.findtext("expect", "").strip(),
                    "dependencies": tool_elem.findtext("dependencies", "").strip(),
                    "time": tool_elem.findtext("time", "").strip(),
                    "failure_policy": tool_elem.findtext("failure_policy", "request user").strip()
                }
                params_elem = tool_elem.find("params")
                if params_elem is not None:
                    for param in params_elem:
                        tool_call["params"][param.tag] = param.text.strip()
                parsed_data["tool_calls"].append(tool_call)

        # Find memory queries
        memory_match = re.search(r"<memory>(.*?)</memory>", llm_output, re.DOTALL)
        if memory_match:
            # Wrap in a root tag to ensure it's valid XML
            memory_xml = f"<root>{memory_match.group(1)}</root>"
            memory_root = ET.fromstring(memory_xml)
            for child in memory_root:
                if child.tag.startswith("search_"):
                    parsed_data["memory_queries"].append(child.text.strip())

    except ET.ParseError as e:
        logger.error(f"Error parsing tool_calls or memory XML: {e}. LLM output fragment might be malformed.")
        # Don't treat the whole output as a reply here, as other parts might have been parsed correctly.
        # The main loop should handle cases where no action is found.

    # If after all parsing, no action is found, it's an issue the main loop should handle.
    # It might decide to fallback to a reply or ask for clarification.
    if not (parsed_data["reply"] or parsed_data["clarify"] or parsed_data["all_plan"] or parsed_data["tool_calls"]):
         logger.warning(f"LLM output could not be parsed into a specific action. Fallback needed. Output: {llm_output}")
         # The calling function will be responsible for the fallback
         parsed_data["reply"] = llm_output # Provide a fallback reply

    return parsed_data
